"""
Agent 后端服务 - 基于 FastAPI 提供 API 接口（简化版）

提供同步处理模式的代理服务，直接返回结果或中断数据，
简化了状态管理和通信流程
"""
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import Dict, Any, Optional, List
import json
import uuid

from langgraph.checkpoint.memory import InMemorySaver
from langgraph.types import interrupt, Command
from langgraph.prebuilt import create_react_agent
from langchain_openai import ChatOpenAI
from langchain_tavily import TavilySearch
import os
import uvicorn

# 会话存储 - 保存每个用户的智能体实例和状态
# 结构: {user_id: 
#           {"agent": agent, 
#            "session_id": session_id}
#       }
sessions = {}

# 创建 LLM
model = ChatOpenAI(model="gpt-4o-mini")

# 定义生命周期上下文管理器
from contextlib import asynccontextmanager
from langgraph.checkpoint.postgres import PostgresSaver

@asynccontextmanager
async def lifespan(app: FastAPI):
    
    # 初始化 PostgreSQL 检查点保存器
    with PostgresSaver.from_conn_string("postgresql://postgres:yourpassword@localhost/postgres?sslmode=disable") as checkpointer:
        checkpointer.setup()

        # 将检查点保存器设置为全局可访问
        app.state.checkpointer = checkpointer

        yield

# 使用生命周期上下文管理器初始化应用
app = FastAPI(
    title="Agent 后端服务", 
    description="基于 LangGraph 提供 AI Agent 服务",
    lifespan=lifespan
)

# 客户端发起的智能体请求
class AgentRequest(BaseModel):
    user_id: str
    query: str
    system_message: Optional[str] = "你会使用工具来帮助用户。如果工具使用被拒绝，请提示用户。"

# 智能体给予的响应
class AgentResponse(BaseModel):

    # 一次会话的唯一标识符
    session_id: str

    # 三个状态：interrupted, completed, error
    status: str  

    #error时的提示消息
    message: Optional[str] = None

    #completed时的结果消息
    result: Optional[Dict[str, Any]] = None

    #interrupted时的中断消息
    interrupt_data: Optional[Dict[str, Any]] = None

# 客户端给予的反馈响应
class InterruptResponse(BaseModel):
    user_id: str
    session_id: str

    # 响应类型：accept, reject, edit
    response_type: str  

    # 如果是edit类型，可能需要额外的参数
    args: Optional[Dict[str, Any]] = None

# 系统信息响应
class SystemInfoResponse(BaseModel):
    sessions_count: int
    active_users: List[str]
    
# 处理智能体返回结果：可能是中断，也可能是最终结果
def process_agent_result(
    session_id: str, 
    result: Dict[str, Any]
) -> AgentResponse:
    """
    处理智能体执行结果，统一处理中断和结果
    
    Args:
        result: 智能体执行结果
        session_id: 会话ID
        
    Returns:
        AgentResponse: 标准化的响应对象
    """
    try:
        # 检查是否有中断
        if "__interrupt__" in result:
            interrupt_data = result["__interrupt__"][0].value
            
            # 返回中断信息
            return AgentResponse(
                session_id=session_id,
                status="interrupted",
                interrupt_data=interrupt_data
            )
        
        # 如果没有中断，返回最终结果
        return AgentResponse(
            session_id=session_id,
            status="completed",
            result=result
        )
    
    except Exception as e:
        return AgentResponse(
            session_id=session_id,
            status="error",
            message=f"处理智能体结果时出错: {str(e)}"
        )

# Tavily搜索工具，需要人工审核/批准
def tavily_search(query: str, search_depth: Optional[str] = "basic"):
    """
    使用Tavily进行网络搜索
    
    参数:
    - query: 搜索查询
    - search_depth: 搜索深度，可选值为 "basic"（基础搜索）或 "advanced"（深度搜索）
    """
    # 中断执行，等待人工审核
    response = interrupt({
        "tool": "tavily_search",
        "args": {
            "query": query,
            "search_depth": search_depth
        },
        "message": f"准备使用Tavily搜索：\n- 查询内容: {query}\n- 搜索深度: {search_depth}\n\n是否允许继续？\n输入 'yes' 接受，'no' 拒绝，或 'edit' 修改查询关键词",
    })   

    # 处理人工响应
    if response["type"] == "accept":
        pass
    elif response["type"] == "edit":
        query = response["args"]["query"]
    else:
        return f"该工具被拒绝使用，请尝试其他方法或拒绝回答问题。"
    
    # 使用LangChain的TavilySearch工具进行实际搜索
    try:
        search = TavilySearch(api_key=os.environ.get("TAVILY_API_KEY"))
        search_results = search.invoke({
            "query": query,
            "search_depth": search_depth,
            "max_results": 5
        })
    except Exception as e:
       return f"该工具调用失败: {str(e)}"
    
    return json.dumps(search_results, ensure_ascii=False, indent=2)

#创建智能体：React Agent
def create_tavily_search_agent(user_id: str):

    agent = create_react_agent(
        model=model,
        tools=[tavily_search],
        checkpointer=app.state.checkpointer,  # 使用全局检查点保存器
    )
    
    return agent

# 开放接口：创建智能体并调用
# 注意：此接口为同步版本，直接返回结果或中断数据
@app.post("/agent/invoke", response_model=AgentResponse)
def invoke_agent(request: AgentRequest):
    """启动智能体处理用户请求 - 同步版本，等待执行完成或中断"""
    user_id = request.user_id
    
    # 检查是否已有会话，如果没有则创建新会话
    if user_id not in sessions:

        # 只在创建新会话时生成新的会话ID
        session_id = str(uuid.uuid4())
        agent = create_tavily_search_agent(user_id)
        sessions[user_id] = {
            "agent": agent,
            "session_id": session_id
        }
    else:
        # 使用现有会话的ID和agent
        agent = sessions[user_id]["agent"]
        session_id = sessions[user_id]["session_id"]
        
    # 初始化智能体输入
    messages = [
        {"role": "system", "content": request.system_message},
        {"role": "user", "content": request.query}
    ]
    
    try:

        # 先调用智能体
        result = agent.invoke({"messages": messages},config={"configurable": {"thread_id": session_id}})

        # 再处理结果
        return process_agent_result(session_id, result)
    
    except Exception as e:
        # 异常处理
        return AgentResponse(
            session_id=session_id,
            status="error",
            message=f"处理请求时出错: {str(e)}"
        )

# 开放接口：恢复被中断的智能体执行
# 注意：此接口为同步版本，等待执行完成或再次中断
@app.post("/agent/resume", response_model=AgentResponse)
def resume_agent(response: InterruptResponse):
    """恢复被中断的智能体执行 - 同步版本，等待执行完成或再次中断"""
    user_id = response.user_id
    client_session_id = response.session_id
    
    # 检查用户会话是否存在
    if user_id not in sessions:
        raise HTTPException(status_code=404, detail=f"用户会话 {user_id} 不存在")
    
    # 检查会话ID是否匹配
    server_session_id = sessions[user_id]["session_id"]
    if server_session_id != client_session_id:
        raise HTTPException(status_code=400, detail="会话ID不匹配，可能是过期的请求")
    
    # 获取智能体和配置
    agent = sessions[user_id]["agent"]
    
    # 构造响应数据
    command_data = {
        "type": response.response_type
    }
    
    # 如果提供了参数，添加到响应数据中
    if response.args:
        command_data["args"] = response.args
    
    try:
        # 先恢复智能体执行
        result = agent.invoke(Command(resume=command_data), config={"configurable": {"thread_id": server_session_id}})

        # 再处理结果
        return process_agent_result(server_session_id, result)
    
    except Exception as e:
        # 异常处理
        return AgentResponse(
            session_id=server_session_id,
            status="error",
            message=f"恢复执行时出错: {str(e)}"
        )

####################################################
# 以下为其他辅助接口
####################################################
@app.get("/agent/status/{user_id}")
def get_agent_status(user_id: str):
    """获取用户智能体的会话状态"""
    if user_id not in sessions:
        return {
            "status": "not_found",
            "message": f"用户 {user_id} 的会话不存在"
        }
    
    session = sessions[user_id]
    return {
        "user_id": user_id,
        "session_id": session["session_id"],
        "status": "active"
    }

@app.get("/system/info", response_model=SystemInfoResponse)
def get_system_info():
    """获取系统状态信息"""
    return SystemInfoResponse(
        sessions_count=len(sessions),
        active_users=list(sessions.keys())
    )

@app.delete("/agent/session/{user_id}")
def delete_agent_session(user_id: str):
    """删除用户会话"""
    if user_id not in sessions:
        raise HTTPException(status_code=404, detail=f"用户 {user_id} 的会话不存在")
    
    # 删除会话
    del sessions[user_id]
    
    return {
        "status": "success",
        "message": f"用户 {user_id} 的会话已删除"
    }

# 启动服务器
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8001)
